#!/usr/bin/python
#Author: Jimmy Saw
#Usage: python get_lineage.py blastfile.tbl 50
#50 is max number of top hits to return
#It is supposed to give you the parent phylum of the hit organism from a BLAST result
#Needs to be in tabular format parsed by Blast2table Perl script

import sys
import re
from Bio import Entrez

Entrez.email = "jimmysaw@gmail.com"

blastfile = sys.argv[1]
bf = open(blastfile, "rU")
bfl = bf.readlines()

maxorgs = sys.argv[2]

hit1 = re.compile('\[(\w+ \w+ \w+)\].*')
hit2 = re.compile('\[(.*)\]')
lin = re.compile('(\w+ \w+).*')
#lin = re.compile('(\w+).*')
acc = re.compile('\w+|(.*)|')

num = len(bfl)

max = 0

if num > int(maxorgs):
    max = int(maxorgs)
else:
    max = num

for line in bfl[0:max]:
    s = line.split('\t')
    evalue = s[1]
    identities = s[3]
    similarities = s[4]
    locus_tag = s[5]
    accession = s[9]
    desc = s[13].rstrip()
    searchterm = ""
    if hit1.search(desc):
        p = hit1.search(desc)
        searchterm = p.group(1)
    elif hit2.search(desc):
        p = hit2.search(desc)
        searchterm = p.group(1)
    else:
        break

    #pattern = rx1.match(searchterm)
    handle = Entrez.esearch(db="Taxonomy", term=searchterm)
    record = Entrez.read(handle)
    taxid = record["IdList"]
    if len(taxid) != 0:
        h = Entrez.efetch(db="Taxonomy", id=taxid, retmode="xml")
        rec = Entrez.parse(h)
        for r in rec:
            lineage = r.get("Lineage")
            l = lineage.split(';')
            #print l[3].lstrip()
            #print "Search term used =", searchterm #pattern.group(1)
            #print "Lineage found =", lineage
            print locus_tag, "Phylum =", l[2].lstrip()
    else:
        print locus_tag, "Can't query NCBI Taxonomy. Taxid not found!"
        print "Search term used =", searchterm

bf.close()
